import h5py
import os
import numpy as np
import SimpleITK as sitk
from tqdm import tqdm
import glob
# 四种模态的mri图像
modalities = ('flair', 't1ce', 't1', 't2')

# train
train_set = {
        'root': '/home/yinsiyuan/paper/BraTS2015/data',  # 四个模态数据所在地址
        'out': '/home/yinsiyuan/paper/BraTS2015/dataset/',  # 预处理输出地址
        'flist': '/home/yinsiyuan/paper/BraTS2015/data/train.txt',  # 训练集名单（有标签），如果为None则自动扫描所有病例
        }


def process_h5(case_path, out_path, case_name=None):
    """ Save the data with dtype=float32.
        z-score is used but keep the background with zero! 
    
    Args:
        case_path: 病例文件夹路径（包含所有模态文件的文件夹）
        out_path: 输出路径
        case_name: 病例名称，如果为None则从路径自动提取
    """
    # 自动检测路径格式（适配不同的下载方式）
    if case_name is None:
        # 从路径提取病例名称
        case_name = os.path.basename(case_path.rstrip(os.sep))
    
    # 构建文件路径 - 支持两种命名方式
    # 方式1: BraTS2015_XXXXX_flair.nii.gz (Kaggle格式)
    # 方式2: BraTS2015_XXXXX_flair.nii.gz (原格式)
    base_name = case_name
    if not base_name.startswith('BraTS2015_'):
        # 如果路径是文件夹内的，尝试从父目录获取
        parent_name = os.path.basename(os.path.dirname(case_path))
        if parent_name.startswith('BraTS2015_'):
            base_name = parent_name
    
    # 尝试不同的路径格式
    label_file = None
    image_files = []
    
    # 方式1: 如果case_path是包含所有文件的文件夹
    if os.path.isdir(case_path):
        # 在文件夹内查找文件
        seg_files = glob.glob(os.path.join(case_path, '*seg*.nii.gz'))
        if seg_files:
            label_file = seg_files[0]
        
        for modal in modalities:
            modal_files = glob.glob(os.path.join(case_path, f'*{modal}*.nii.gz'))
            if modal_files:
                image_files.append(modal_files[0])
    else:
        # 方式2: 传统格式 - case_path是文件前缀路径
        label_file = case_path + 'seg.nii.gz'
        image_files = [case_path + modal + '.nii.gz' for modal in modalities]
    
    # 检查文件是否存在
    if label_file is None or not os.path.exists(label_file):
        raise FileNotFoundError(f"标签文件未找到: {label_file}")
    
    for img_file in image_files:
        if not os.path.exists(img_file):
            raise FileNotFoundError(f"图像文件未找到: {img_file}")
    
    # SimpleITK读取图像默认是 DxHxW，这里转为 HxWxD
    label = sitk.GetArrayFromImage(sitk.ReadImage(label_file)).transpose(1,2,0)
    print(f"标签形状: {label.shape}")
    
    # 堆叠四种模态的图像，4 x (H,W,D) -> (4,H,W,D) 
    images = np.stack([sitk.GetArrayFromImage(sitk.ReadImage(img_file)).transpose(1,2,0) for img_file in image_files], 0)  # [4,240,240,155]
    
    # 数据类型转换
    label = label.astype(np.uint8)
    images = images.astype(np.float32)
    
    # 确保输出目录存在
    os.makedirs(out_path, exist_ok=True)
    
    # 构建输出文件名
    output = os.path.join(out_path, case_name + '_mri_norm2.h5')
    # 对第一个通道求和，如果四个模态都为0，则标记为背景(False)
    mask = images.sum(0) > 0
    for k in range(4):

        x = images[k,...]  #
        y = x[mask]

        # 对背景外的区域进行归一化
        x[mask] -= y.mean()
        x[mask] /= y.std()

        images[k,...] = x
    print(case_name,images.shape,label.shape)
    f = h5py.File(output, 'w')
    f.create_dataset('image', data=images, compression="gzip")
    f.create_dataset('label', data=label, compression="gzip")
    f.close()


def doit(dset):
    """
    处理数据集
    
    Args:
        dset: 数据集配置字典，包含 'root', 'out', 'flist' 键
    """
    root = dset['root']
    out_path = dset['out']
    
    # 确保输出目录存在
    os.makedirs(out_path, exist_ok=True)
    
    # 如果有文件列表，使用文件列表；否则自动扫描所有病例文件夹
    if dset.get('flist') and os.path.exists(dset['flist']):
        # 使用文件列表
        file_list = dset['flist']
        subjects = open(file_list, 'r').read().splitlines()
        subjects = [s.strip() for s in subjects if s.strip()]  # 去除空行和空白
        
        # 处理文件名格式
        names = []
        for sub in subjects:
            if sub.startswith('BraTS2015_'):
                names.append(sub)
            else:
                names.append('BraTS2015_' + sub)
        
        # 构建路径 - 适配不同的数据组织结构
        paths = []
        for name in names:
            # 尝试多种路径格式
            path1 = os.path.join(root, name, name + '_')  # 传统格式
            path2 = os.path.join(root, name)  # kagglehub格式（直接是文件夹）
            
            if os.path.exists(path2) and os.path.isdir(path2):
                paths.append((path2, name))
            elif os.path.exists(os.path.dirname(path1)) or os.path.exists(os.path.dirname(path2)):
                # 如果父目录存在，尝试查找
                parent_dir = os.path.dirname(path2) if os.path.exists(os.path.dirname(path2)) else os.path.dirname(path1)
                for item in os.listdir(parent_dir):
                    if item.startswith(name):
                        full_path = os.path.join(parent_dir, item)
                        if os.path.isdir(full_path):
                            paths.append((full_path, name))
                            break
    else:
        # 自动扫描所有病例文件夹
        print("未找到文件列表，自动扫描数据集目录...")
        if not os.path.exists(root):
            raise FileNotFoundError(f"数据根目录不存在: {root}")
        
        # 查找所有以 BraTS2015_ 开头的文件夹
        all_items = os.listdir(root)
        case_folders = [item for item in all_items if item.startswith('BraTS2015_') and os.path.isdir(os.path.join(root, item))]
        
        if not case_folders:
            # 如果在子目录中
            for item in all_items:
                sub_path = os.path.join(root, item)
                if os.path.isdir(sub_path):
                    sub_items = os.listdir(sub_path)
                    case_folders.extend([os.path.join(item, ci) for ci in sub_items if ci.startswith('BraTS2015_') and os.path.isdir(os.path.join(sub_path, ci))])
        
        case_folders.sort()
        names = [os.path.basename(cf) if os.sep in cf else cf for cf in case_folders]
        paths = [(os.path.join(root, cf), name) for cf, name in zip(case_folders, names)]
        
        print(f"自动检测到 {len(paths)} 个病例")
    
    # 处理每个病例
    failed_cases = []
    for path, case_name in tqdm(paths, desc="处理病例"):
        try:
            process_h5(path, out_path, case_name)
        except Exception as e:
            print(f"\n处理病例 {case_name} 失败: {str(e)}")
            failed_cases.append((case_name, str(e)))
            continue
    
    print(f'\n处理完成！成功: {len(paths) - len(failed_cases)}, 失败: {len(failed_cases)}')
    if failed_cases:
        print("\n失败的病例:")
        for case_name, error in failed_cases:
            print(f"  {case_name}: {error}")


if __name__ == '__main__':
    doit(train_set)
